from matplotlib import pyplot as plt
import numpy as np
from openpyxl import load_workbook

# x = ['bias_add', 'avg_pool', 'softmax', 'conv2d', 'batch_norm', 'max_pool',
#      'relu', 'reduce_mean', 'reduce_max', 'sigmoid', 'tanh', 'matmul', 'reduce_min', 'reduce_sum', 'abs', 'square']
x = ['op1', 'op2', 'op3', 'op4', 'op5', 'op6', 'op7', 'op8', 'op9', 'op10',
     'op11', 'op12']

# y_hard_mre = [0.30, 7.52, 1.61, 6.09, 4.66, 6.51, 4.77, 4.77, 2.06, 0.10, 0.10, 7.97]
# y_hard_mare = [0, 10.95, 1.95, 10.69, 5.04, 6.24, 3.49, 3.73, 2.50, 0.37, 0.17, 5.71]
# y_hard_mre_weighted = [13.84, 43.71, 54.27, 7.50, 44.80, 72.04, 79.80, 7.53, 74.52, 31.03, 98.22, 8.68]
# y_hard_mare_weighted = [29.57, 40.52, 52.07, 11.54, 19.77, 72.56, 72.66, 5.62, 85.88, 47.11, 98.78, 6.12]
#
# y_hard_mre_weighted = [x * 150 for x in y_hard_mre_weighted]
# y_hard_mre = [x * 150 for x in y_hard_mre]
# y_hard_mare_weighted = [x * 150 for x in y_hard_mare_weighted]
# y_hard_mare = [x * 150 for x in y_hard_mare]
#
# y_hard = [(y_hard_mre[i] + y_hard_mare[i]) / 2 for i in range(12)]
# y_hard_weighted = [(y_hard_mre_weighted[i] + y_hard_mare_weighted[i]) / 2 for i in range(12)]
#
# y_medium_mre = [42.70, 46.12, 1.95, 47.97, 32.28, 36.25, 40.07, 39.67, 32.67, 0.1, 0.1, 41.29]
# y_medium_mare = [13.87, 42.78, 1.95, 42.57, 39.79, 35.15, 36.84, 46.72, 35.0, 0.37, 0.17, 32.38]
# y_medium_mre_weighted = [54.08, 53.99, 54.27, 69.81, 32.28, 36.25, 67.04, 46.17, 90.60, 31.03, 97.88, 47.71]
# y_medium_mare_weighted = [64.55, 50.94, 52.07, 47.11, 43.96, 64.95, 68.85, 50.54, 97.87, 47.11, 98.78, 35.86]
#
# y_medium_mre_weighted = [x * 150 for x in y_medium_mre_weighted]
# y_medium_mre = [x * 150 for x in y_medium_mre]
# y_medium_mare_weighted = [x * 150 for x in y_medium_mare_weighted]
# y_medium_mare = [x * 150 for x in y_medium_mare]
#
# y_medium = [(y_medium_mre[i] + y_medium_mare[i]) / 2 for i in range(12)]
# y_medium_weighted = [(y_medium_mre_weighted[i] + y_medium_mare_weighted[i]) / 2 for i in range(12)]
#
# y = [int((y_medium[i] + y_hard[i]) / 2) for i in range(12)]
# y_weighted = [int((y_medium_weighted[i] + y_hard_weighted[i]) / 2) for i in range(12)]
# print(y)
# print(y_weighted)

bar_width = 0.4

wb = load_workbook('./数据生成算法实验.xlsx')
sheet = wb['带权重采样']

tf_hard = []
torch_hard = []
mnn_hard = []
tf_medium = []
torch_medium = []
mnn_medium = []
tf_hard_w = []
torch_hard_w = []
mnn_hard_w = []
tf_medium_w = []
torch_medium_w = []
mnn_medium_w = []

tf_hard_mare = []
torch_hard_mare = []
mnn_hard_mare = []
tf_medium_mare = []
torch_medium_mare = []
mnn_medium_mare = []
tf_hard_w_mare = []
torch_hard_w_mare = []
mnn_hard_w_mare = []
tf_medium_w_mare = []
torch_medium_w_mare = []
mnn_medium_w_mare = []
for r in range(2, 139):
    name = sheet.cell(row=r, column=1).value
    if name is not None:
        special = name.find('softmax') * name.find('sigmoid') * name.find('tanh')
        forced = sheet.cell(row=r, column=11).value != '/'
        pattern = sheet.cell(row=r, column=3).value
        if name.find('tf') != -1:
            r_random = sheet.cell(row=r, column=6).value
            r_w = sheet.cell(row=r, column=11).value if forced else sheet.cell(row=r, column=5).value
            if special >= 0:
                if pattern == 'MRE':
                    tf_hard.append(r_random)
                    tf_medium.append(r_random)
                    tf_hard_w.append(r_w)
                    tf_medium_w.append(r_w)
                else:
                    tf_hard_mare.append(r_random)
                    tf_medium_mare.append(r_random)
                    tf_hard_w_mare.append(r_w)
                    tf_medium_w_mare.append(r_w)
            else:
                if r_random <= 0.15 and r_random != 0.108:
                    if pattern == 'MRE':
                        tf_hard.append(r_random)
                        tf_hard_w.append(r_w)
                    else:
                        tf_hard_mare.append(r_random)
                        tf_hard_w_mare.append(r_w)
                else:
                    if pattern == 'MRE':
                        tf_medium.append(r_random)
                        tf_medium_w.append(r_w)
                    else:
                        tf_medium_mare.append(r_random)
                        tf_medium_w_mare.append(r_w)
        elif name.find('torch') != -1 or name.find('F') != -1:
            r_random = sheet.cell(row=r, column=6).value
            r_w = sheet.cell(row=r, column=11).value if forced else sheet.cell(row=r, column=5).value
            if special >= 0:
                if pattern == 'MRE':
                    torch_hard.append(r_random)
                    torch_medium.append(r_random)
                    torch_hard_w.append(r_w)
                    torch_medium_w.append(r_w)
                else:
                    torch_hard_mare.append(r_random)
                    torch_medium_mare.append(r_random)
                    torch_hard_w_mare.append(r_w)
                    torch_medium_w_mare.append(r_w)
            else:
                if r_random <= 0.15 and r_random != 0.108:
                    if pattern == 'MRE':
                        torch_hard.append(r_random)
                        torch_hard_w.append(r_w)
                    else:
                        torch_hard_mare.append(r_random)
                        torch_hard_w_mare.append(r_w)
                else:
                    if pattern == 'MRE':
                        torch_medium.append(r_random)
                        torch_medium_w.append(r_w)
                    else:
                        torch_medium_mare.append(r_random)
                        torch_medium_w_mare.append(r_w)
        else:
            r_random = sheet.cell(row=r, column=6).value
            r_w = sheet.cell(row=r, column=11).value if forced else sheet.cell(row=r, column=5).value
            if special >= 0:
                if pattern == 'MRE':
                    mnn_hard.append(r_random)
                    mnn_medium.append(r_random)
                    mnn_hard_w.append(r_w)
                    mnn_medium_w.append(r_w)
                else:
                    mnn_hard_mare.append(r_random)
                    mnn_medium_mare.append(r_random)
                    mnn_hard_w_mare.append(r_w)
                    mnn_medium_w_mare.append(r_w)
            else:
                if r_random <= 0.15 and r_random != 0.108:
                    if pattern == 'MRE':
                        mnn_hard.append(r_random)
                        mnn_hard_w.append(r_w)
                    else:
                        mnn_hard_mare.append(r_random)
                        mnn_hard_w_mare.append(r_w)
                else:
                    if pattern == 'MRE':
                        mnn_medium.append(r_random)
                        mnn_medium_w.append(r_w)
                    else:
                        mnn_medium_mare.append(r_random)
                        mnn_medium_w_mare.append(r_w)

print(len(tf_hard_mare))
print(len(tf_medium_mare))
print(len(tf_medium_w_mare))
print(len(tf_hard_w_mare))
print(len(torch_hard_mare))
print(len(torch_medium_mare))
print(len(torch_medium_w_mare))
print(len(torch_hard_w_mare))
print(len(mnn_hard_mare))
print(len(mnn_medium_mare))
print(len(mnn_medium_w_mare))
print(len(mnn_hard_w_mare))
tf = [int((tf_hard[i] + tf_medium[i]) * 7500) for i in range(12)]
torch = [int((torch_hard[i] + torch_medium[i]) * 7500) for i in range(12)]
mnn = [int((mnn_hard[i] + mnn_medium[i]) * 7500) for i in range(12)]
tf_w = [int((tf_hard_w[i] + tf_medium_w[i]) * 7500) for i in range(12)]
torch_w = [int((torch_hard_w[i] + torch_medium_w[i]) * 7500) for i in range(12)]
mnn_w = [int((mnn_hard_w[i] + mnn_medium_w[i]) * 7500) for i in range(12)]


tf.append(40)
tf.append(0)
tf_w.append(7412)
tf_w.append(6213)
torch.append(40)
torch.append(3095)
torch_w.append(7412)
torch_w.append(4514)
mnn.append(40)
mnn.append(3095)
mnn_w.append(7412)
mnn_w.append(4514)

tf.append(511)
tf.append(0)
tf_w.append(4859)
tf_w.append(15000)
torch.append(511)
torch.append(724)
torch_w.append(4859)
torch_w.append(4450)
mnn.append(511)
mnn.append(724)
mnn_w.append(4859)
mnn_w.append(4450)

print(tf)
print(torch)
print(mnn)
print(tf_w)
print(torch_w)
print(mnn_w)


tf = tf[0: 2] + tf[3: 9] + tf[11: 13] + tf[14:]
tf_w = tf_w[0: 2] + tf_w[3: 9] + tf_w[11: 13] + tf_w[14:]
torch = torch[0: 2] + torch[3: 9] + torch[11: 13] + torch[14:]
torch_w = torch_w[0: 2] + torch_w[3: 9] + torch_w[11: 13] + torch_w[14:]
mnn = mnn[0: 2] + mnn[3: 9] + mnn[11: 13] + mnn[14:]
mnn_w = mnn_w[0: 2] + mnn_w[3: 9] + mnn_w[11: 13] + mnn_w[14:]


# analysis = np.asarray(weights)

# plt.title('Weighted vs. Random Sampling')
# plt.title('Weighted vs. Random Sampling')

plt.barh(y=range(len(x)), width=0, label='$TF\ D_{pm}$', height=bar_width, color='#9467bd')
plt.barh(y=np.arange(len(x))+bar_width, width=0, label='$TF\ D_{r}$', height=bar_width, color='#1f77b4')

plt.barh(y=range(len(x)), width=0, label='$PyTorch\ D_{pm}$', height=bar_width, color='#8c564b')
plt.barh(y=np.arange(len(x))+bar_width, width=0, label='$PyTorch\ D_{r}$', height=bar_width, color='#ff7f0e')

plt.barh(y=range(len(x)), width=0, label='$MNN\ D_{pm}$', height=bar_width, color='#d62728')
plt.barh(y=np.arange(len(x))+bar_width, width=0, label='$MNN\ D_{r}$', height=bar_width, color='#2ca02c')

y = np.asarray([tf, torch, mnn])
y_weighted = np.asarray([tf_w, torch_w, mnn_w])

max_i = np.argmax(y, 0)
min_i = np.argmin(y, 0)
for i in range(12):
    if max_i[i] == 0:
        plt.barh(y=np.asarray([i]) + bar_width, width=[tf[i]], height=bar_width, color='#1f77b4')
    elif max_i[i] == 1:
        plt.barh(y=np.asarray([i]) + bar_width, width=[torch[i]], height=bar_width, color='#ff7f0e')
    else:
        plt.barh(y=np.asarray([i]) + bar_width, width=[mnn[i]], height=bar_width, color='#2ca02c')

    if 3 - max_i[i] - min_i[i] == 0:
        plt.barh(y=np.asarray([i]) + bar_width, width=[tf[i]], height=bar_width, color='#1f77b4')
    elif 3 - max_i[i] - min_i[i] == 1:
        plt.barh(y=np.asarray([i]) + bar_width, width=[torch[i]], height=bar_width, color='#ff7f0e')
    else:
        plt.barh(y=np.asarray([i]) + bar_width, width=[mnn[i]], height=bar_width, color='#2ca02c')

    if min_i[i] == 0:
        plt.barh(y=np.asarray([i]) + bar_width, width=[tf[i]], height=bar_width, color='#1f77b4')
    elif min_i[i] == 1:
        plt.barh(y=np.asarray([i]) + bar_width, width=[torch[i]], height=bar_width, color='#ff7f0e')
    else:
        plt.barh(y=np.asarray([i]) + bar_width, width=[mnn[i]], height=bar_width, color='#2ca02c')

    if y[max_i[i]][i] >= 14000:
        plt.text(y[max_i[i]][i] - 700, i + bar_width / 2 - 0.02, '{0}'.format(y[max_i[i]][i]), ha='center', va='bottom', fontsize=7)
    else:
        plt.text(y[max_i[i]][i] + 800, i + bar_width / 2 - 0.05, '{0}'.format(y[max_i[i]][i]), ha='center', va='bottom', fontsize=9)


max_i = np.argmax(y_weighted, 0)
min_i = np.argmin(y_weighted, 0)
for i in range(12):
    if max_i[i] == 0:
        plt.barh(y=[i], width=[tf_w[i]], height=bar_width, color='#9467bd')
    elif max_i[i] == 1:
        plt.barh(y=[i], width=[torch_w[i]], height=bar_width, color='#8c564b')
    else:
        plt.barh(y=[i], width=[mnn_w[i]], height=bar_width, color='#d62728')

    if 3 - max_i[i] - min_i[i] == 0:
        plt.barh(y=[i], width=[tf_w[i]], height=bar_width, color='#9467bd')
    elif 3 - max_i[i] - min_i[i] == 1:
        plt.barh(y=[i], width=[torch_w[i]], height=bar_width, color='#8c564b')
    else:
        plt.barh(y=[i], width=[mnn_w[i]], height=bar_width, color='#d62728')

    if min_i[i] == 0:
        plt.barh(y=[i], width=[tf_w[i]], height=bar_width, color='#9467bd')
    elif min_i[i] == 1:
        plt.barh(y=[i], width=[torch_w[i]], height=bar_width, color='#8c564b')
    else:
        plt.barh(y=[i], width=[mnn_w[i]], height=bar_width, color='#d62728')

    if y_weighted[max_i[i]][i] >= 14000:
        plt.text(y_weighted[max_i[i]][i] - 700, i - bar_width / 2 - 0.02, '{0}'.format(y_weighted[max_i[i]][i]), ha='center', va='bottom', fontsize=7)
    else:
        plt.text(y_weighted[max_i[i]][i] + 800, i - bar_width / 2 - 0.05, '{0}'.format(y_weighted[max_i[i]][i]), ha='center', va='bottom', fontsize=9)

# for i in range(3):
#     for y_, x_ in enumerate(y_weighted[i]):
#         if x_ >= 14000:
#             plt.text(x_ - 800, y_ - bar_width / 2, '{0}'.format(x_), ha='center', va='bottom', fontsize=7)
#         else:
#             plt.text(x_ + 800, y_ - bar_width / 2, '{0}'.format(x_), ha='center', va='bottom', fontsize=7)

    # for y_, x_ in enumerate(y[i]):
    #     if x_ >= 14000:
    #         plt.text(x_ - 800, y_ + bar_width / 2, '{0}'.format(x_), ha='center', va='bottom', fontsize=7)
    #     else:
    #         plt.text(x_ + 800, y_ + bar_width / 2, '{0}'.format(x_), ha='center', va='bottom', fontsize=7)
#
plt.xlabel("Successful sample(s)", fontsize=11)
plt.ylabel("Op(s)", fontsize=11, labelpad=7)
plt.yticks(np.arange(len(x))+bar_width/2, x, fontsize=11)
plt.xticks(fontsize=11)
plt.legend(loc='lower right', fontsize=7.3)

# plt.rcParams['figure.figsize'] = (28, 21)
plt.savefig('../time/plott/comp.eps', format='eps', dpi=1000)
plt.show()